import os
import pandas as pd
import numpy as np
import pickle
import logging
from tqdm import tqdm
from rdkit import Chem
import xgboost as xgb
from app.features import main_pipeline
from fastapi import FastAPI
from typing import Optional
import uvicorn
from pydantic import BaseModel, Field
from app.molecular_water_solubility_prediction import predict_water_solubility 

class RedoxRequest(BaseModel):
    O_smile: str = Field(..., 
        min_length=1, 
        example="C1C(=O)COC2=NC=CN21",
        description="氧化态分子SMILES（需通过RDKit验证）")
    
    H_smile: str = Field(...,
        min_length=1,
        example="OC1COc2nccn2C1",
        description="还原态分子SMILES（需通过RDKit验证）")

class WaterSolubilityRequest(BaseModel):
    smile: str = Field(..., 
        min_length=1, 
        example="C1C(=O)COC2=NC=CN21",
        description="分子SMILES（需通过RDKit验证）")

app = FastAPI(debug=True)

# 预加载模型组件（避免每次请求重复加载）
artifacts = {
    "model": None,
    "scaler": None,
    "features": []
    }
@app.on_event("startup")
def load_artifacts(model_dir="./app/models"):
    """加载模型相关组件"""
    global artifacts
    try:
        # 获取最新模型版本
        model_files = sorted([f for f in os.listdir(model_dir) if f.endswith("_model.json")])
        if not model_files:
            raise FileNotFoundError("未找到模型文件")
        latest_model = model_files[-1]
        prefix = latest_model.split("_model.json")[0]
        
        # 加载组件
        # with open(os.path.join(model_dir, latest_model), "rb") as f:
        #     artifacts["model"] = pickle.load(f)
        artifacts["model"] = xgb.Booster()
        artifacts["model"].load_model("./app/models/XGBoost_O-H_mordred_model.json")
            
        with open(os.path.join(model_dir, f"{prefix}_scaler.pkl"), "rb") as f:
            artifacts["scaler"] = pickle.load(f)
            
        with open(os.path.join(model_dir, f"{prefix}_features.txt"), "r") as f:
            artifacts["features"] = [line.strip() for line in f.readlines()]
            
        # return artifacts, prefix
    
    except Exception as e:
        logging.error(f"组件加载失败: {str(e)}")
        raise

def validate_smiles(smiles):
    """验证SMILES有效性并返回规范格式"""
    try:
        mol = Chem.MolFromSmiles(smiles)
        if not mol:
            raise ValueError(f"无效的SMILES: {smiles}")
        return Chem.MolToSmiles(mol)
    except Exception as e:
        raise ValueError(f"SMILES验证失败: {str(e)}")

def predict_potential(O_smile: str, H_smile: str, model_dir="./app/results"):
    """预测电极电势的核心函数"""
    try:
        # 输入验证与标准化
        O_smile = validate_smiles(O_smile)
        H_smile = validate_smiles(H_smile)
        
        # 构建输入数据框架
        input_data = pd.DataFrame([[O_smile, H_smile]], 
                                columns=['O_smile', 'H_smile'])
        
        # 执行特征工程
        processed_data = main_pipeline(input_data)
        
        # 加载模型组件
        # artifacts, _ = load_artifacts(model_dir)
        
        # 特征对齐检查
        required_features = artifacts["features"]
        if not set(required_features).issubset(processed_data.columns):
            missing = set(required_features) - set(processed_data.columns)
            raise ValueError(f"特征缺失: {missing}")
            
        # 数据标准化
        X_scaled = artifacts["scaler"].transform(processed_data[required_features].values)
        
        # 将 NumPy 数组转换为 DMatrix
        dmatrix_data = xgb.DMatrix(X_scaled)
        
        # 执行预测
        prediction = artifacts["model"].predict(dmatrix_data)[0].item() 
        
        # return round(prediction, 3)  # 保留三位小数
        return prediction
    
    except Exception as e:
        logging.error(f"预测失败: {str(e)}")
        raise

@app.post("/predict/standard_potential")
async def api_predict_standard_potential(request: RedoxRequest):  # 使用模型接收请求体
    try:
        # 通过模型实例访问参数
        O_smile = request.O_smile  
        H_smile = request.H_smile
        logging.info(f"请求参数 - O_smile: {O_smile}, H_smile: {H_smile}")
        
        potential = predict_potential(O_smile, H_smile)
        return {"status": "success", "standard potential": potential}
    
    except ValueError as ve:  # 捕获分子验证异常
        logging.error(f"SMILES格式错误: {str(ve)}")
        return {"status": "error", "message": str(ve)}
    
    except Exception as e:
        logging.error(f"API预测失败: {str(e)}", exc_info=True)
        return {"status": "error", "message": "内部服务错误"}

@app.post("/predict/water_solubility")
async def api_predict_water_solubility(request: WaterSolubilityRequest):
    try:
        # 通过模型实例访问参数
        smile = request.smile 
        logging.info(f"请求参数 - SMILES: {smile}")
        
        water_solubility = predict_water_solubility(smile)

        if water_solubility is None:
            return {"status": "error", "message": f"错误的SMILES {smile}"}
        return {"status": "success", "water solubility": water_solubility}
    
    except ValueError as ve:  # 捕获分子验证异常
        logging.error(f"SMILES格式错误: {str(ve)}")
        return {"status": "error", "message": str(ve)}
    
    except Exception as e:
        logging.error(f"API预测失败: {str(e)}", exc_info=True)
        return {"status": "error", "message": "内部服务错误"}

if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO,
                        format="%(asctim)es - %(levelname)s - %(message)s")
    
    # uvicorn.run(app, host="0.0.0.0", port=8000, workers=4, reload=False)
    